#ifndef PW_PBL_FORCEALIGN_HPP
#define PW_PBL_FORCEALIGN_HPP

#include "Math.hpp"
#include "RustMimic.hpp"

namespace pw_pbl
{
	/// 非对齐类型引用
	/// @note UnalignRef<T>表示实际为非对齐类型，UnalignRef<T&>表示实际为对齐类型但提供非对齐类型兼容的接口，UnalignRef<>提供公共静态操作
	template <typename... T>
	struct UnalignRef;

	/// 强制指定对齐类型
	/// @tparam T 原始类型
	/// @tparam ALIGN 对齐值
	/// @note ForceAlign<> = ForceAlign<void>为公共空基类，可用std::derived_from<ForceAlign<>>检测是否ForceAlign类型
	/// @note ForceAlign<T> 为次级公共空基类，可用std::derived_from<ForceAlign<T>>检测是否ForceAlign<T,ALIGN>类型
	template <typename T = void, size_t... ALIGN> requires(sizeof...(ALIGN) <= 1)
	struct ForceAlign;

	template <typename... T>
	struct MaybeUnalignTraits;

	template <typename T>
	struct MaybeUnalignTraits<T>
	{
		using raw_type = T;
	};

	template <typename T, size_t ALIGN>
	struct MaybeUnalignTraits<ForceAlign<T, ALIGN>>
	{
		using raw_type = T;
	};

	template <>
	struct MaybeUnalignTraits<>
	{
		template <typename T>
		static constexpr T read(const T &ref) { return ref; }

		template <typename T>
		static constexpr void write(T &ref, T v) { ref = std::forward<T>(v); }

		template <typename T>
		static constexpr void alter(T &ref, std::invocable<T &> auto func) { std::move(func)(ref); }

		template <typename T, size_t ALIGN>
		static constexpr T read(const ForceAlign<T, ALIGN> &ref) { return ref.read(); }

		template <typename T, size_t ALIGN>
		static constexpr void write(ForceAlign<T, ALIGN> &ref, T v) { ref.write(std::forward<T>(v)); }

		template <typename T, size_t ALIGN>
		static constexpr void alter(ForceAlign<T, ALIGN> &ref, std::invocable<T &> auto func) { ref.alter(std::move(func)); }
	};

	template <typename T>
	concept ccMaybeUnalignIntegral = std::integral<typename MaybeUnalignTraits<T>::raw_type>;

	template <typename T>
	concept ccMaybeUnalignUnsignedIntegral = std::unsigned_integral<typename MaybeUnalignTraits<T>::raw_type>;

	template <typename T>
	concept ccMaybeUnalignSignedIntegral = std::signed_integral<typename MaybeUnalignTraits<T>::raw_type>;

	template <>
	struct UnalignRef<>
	{
		template <typename T>
		static constexpr T checked_read(const T *ptr)
		{
			if consteval
			{
				return *ptr;
			}
			else
			{
				if constexpr (alignof(T) == 1)
					return *ptr;
				else if (((uintptr_t)ptr & EMask<T>) == 0)
					return *ptr;
				else
				{
					static_assert(std::is_trivially_copyable_v<T>);
					alignas(T) std::byte tmp[sizeof(T)];
					memcpy(tmp, ptr, sizeof(T));
					return reinterpret_cast<const T &>(tmp);
				}
			}
		}

		template <typename T> requires(!std::is_const_v<T>)
		static constexpr void checked_write(T *ptr, T v)
		{
			if consteval
			{
				*ptr = v;
			}
			else
			{
				if constexpr (alignof(T) == 1)
					*ptr = v;
				else if (((uintptr_t)ptr & EMask<T>) == 0)
					*ptr = v;
				else
				{
					static_assert(std::is_trivially_copyable_v<T>);
					memcpy(ptr, &v, sizeof(T));
				}
			}
		}

		template <typename T> requires(!std::is_const_v<T>)
		static constexpr void checked_alter(T *ptr, std::invocable<T &> auto func)
		{
			if consteval
			{
				func(*ptr);
			}
			else
			{
				if constexpr (alignof(T) == 1)
					func(*ptr);
				else if (((uintptr_t)ptr & EMask<T>) == 0)
					func(*ptr);
				else
				{
					static_assert(std::is_trivially_copyable_v<T>);
					alignas(T) std::byte tmp[sizeof(T)];
					memcpy(tmp, ptr, sizeof(T));
					func(reinterpret_cast<T &>(tmp));
					memcpy(ptr, reinterpret_cast<T &>(tmp), sizeof(T));
				}
			}
		}

	private:
		template <typename T>
		static constexpr uintptr_t EMask = ~(std::numeric_limits<uintptr_t>::max() << Math::Log2<alignof(T)>);
		constexpr UnalignRef() = default;
	};

	template <>
	struct ForceAlign<void>
	{
	protected:
		constexpr ForceAlign() = default;
	};

	template <typename T> requires std::same_as<T, std::decay_t<T>> // 必须是纯值类型
	struct ForceAlign<T> : ForceAlign<void>
	{
		using raw_type = T;

	protected:
		constexpr ForceAlign() = default;
	};

	template <typename T, size_t ALIGN> requires(std::has_single_bit(ALIGN))
	struct alignas(ALIGN) ForceAlign<T, ALIGN> : ForceAlign<T>
	{
		static constexpr bool EIsAligned = alignof(T) <= ALIGN;

		constexpr ForceAlign() noexcept { static_assert(is_valid()); }
		constexpr ForceAlign(T v) noexcept : raw(std::move(v)) { static_assert(is_valid()); }

		constexpr UnalignRef<std::conditional_t<EIsAligned, T &, T>> ref() { return {this->raw}; }
		constexpr UnalignRef<std::conditional_t<EIsAligned, const T &, const T>> ref() const { return {this->raw}; }

		constexpr T read() const { return this->raw; }
		constexpr void write(T v) { this->raw = std::move(v); }
		constexpr void alter(std::invocable<T &> auto func)
		{
			if constexpr (EIsAligned)
				std::move(func)(this->raw);
			else
			{
				T tmp = this->raw;
				std::move(func)(tmp);
				this->raw = std::move(tmp);
			}
		}

	private:
		T raw [[gnu::packed]];

		static constexpr bool is_valid() { return alignof(ForceAlign) == ALIGN && sizeof(ForceAlign) == (alignof(T) >= ALIGN ? sizeof(T) : ALIGN); };
	};

	template <typename T> requires(!std::is_reference_v<T>)
	struct UnalignRef<T>
	{
		T *const ptr;

		constexpr UnalignRef(T &v) noexcept : ptr(&v) {}
		constexpr UnalignRef(UnalignRef<T &> v) noexcept : ptr(v.ptr) {}

		constexpr T read() const { return UnalignRef<>::checked_read(ptr); }
		constexpr void write(T v) requires(!std::is_const_v<T>) { UnalignRef<>::checked_write(ptr, std::move(v)); }
		constexpr void alter(std::invocable<T &> auto func) requires(!std::is_const_v<T>) { UnalignRef<>::checked_alter(ptr, std::move(func)); }

		template <typename U> requires(sizeof(U) == sizeof(T) && !std::is_reference_v<U>)
		inline UnalignRef<U> cast_to() const { return {*reinterpret_cast<U *>(ptr)}; }
	};

	template <typename T>
	struct UnalignRef<T &>
	{
		T *const ptr;

		constexpr UnalignRef(T &v) noexcept : ptr(&v) {}

		constexpr T read() const { return *ptr; }
		constexpr void write(T v) requires(!std::is_const_v<T>) { *ptr = std::move(v); }
		constexpr void alter(std::invocable<T &> auto func) requires(!std::is_const_v<T>) { std::move(func)(*ptr); }

		template <typename U> requires(sizeof(U) == sizeof(T) && alignof(U) <= alignof(T))
		inline UnalignRef<U &> cast_to() const { return {*reinterpret_cast<std::remove_reference_t<U> *>(ptr)}; }
	};

	template <typename T>
	class Option<UnalignRef<T>>
	{
		UnalignRef<T> raw;

	public:
		constexpr Option() : raw(Option<T &>{}.get_unchecked()) {}
		constexpr Option(UnalignRef<T> data) : raw(data) {}
		constexpr bool is_none() const { return raw.ptr == nullptr; }
		constexpr decltype(auto) get_unchecked(this auto &&self) { return std::forward_like<decltype(self)>(self.raw); }
	};

} // namespace pw_pbl

#endif // PW_PBL_FORCEALIGN_HPP
